import os
import re
import json
from typing import Dict


# 加载文档结构
def load_document_structure(structure_path: str) -> Dict:
    with open(structure_path, 'r', encoding='utf-8') as f:
        return json.load(f)


# 从文档结构中提取章节标题和摘要信息，用于章节匹配
def extract_sections_info(doc_structure: Dict) -> str:
    sections_info = []

    def traverse_structure(structure, prefix=""):
        for key, data in structure.items():
            full_key = f"{prefix}{key}" if prefix else key
            title = data.get("title", "")
            summary = data.get("summary", "")
            sections_info.append(f"章节 {full_key}: {title}\n摘要: {summary}")

            if "children" in data and data["children"]:
                # new_prefix = f"{full_key}." if full_key else ""
                # traverse_structure(data["children"], new_prefix)
                traverse_structure(data["children"])

    traverse_structure(doc_structure)
    return "\n\n".join(sections_info)


# 递归查找所有层级的代码段信息
def find_code_segments(content):
    code_segments = {}

    # 如果当前层级有代码段，提取它们
    if "代码段" in content:
        code_segments.update(content["代码段"])

    # 查找子章节中的代码段
    if "children" in content:
        for sub_content in content["children"].values():
            code_segments.update(find_code_segments(sub_content))

    return code_segments


# 加载章节内容
def load_full_chapter_content(chapter_id: str, chapters_dir: str) -> Dict:
    main_chapter_id = chapter_id.split('.')[0]

    chapter_path = os.path.join(chapters_dir, f"{main_chapter_id}.json")
    if os.path.exists(chapter_path):
        with open(chapter_path, 'r', encoding='utf-8') as f:
            chapter_content = json.load(f)
    else:
        raise RuntimeError(f"文件 {chapter_path} 不存在！")

    if main_chapter_id not in chapter_content:
        return f"章节 {main_chapter_id} 在文档 {chapter_path} 中不存在"

    is_chapter_16 = main_chapter_id == '16'

    chapter_code_segments = None if is_chapter_16 else find_code_segments(chapter_content[main_chapter_id])

    return chapter_content, chapter_code_segments, is_chapter_16


# 从章节内容中提取指定路径的文本内容，包括代码段
def extract_specific_chapter_text(chapter_content: Dict, chapter_path: str) -> str:
    if not chapter_content:
        return ""

    # 分解章节路径，例如将"2.3.1"分解为['2', '3', '1']
    path_components = chapter_path.split('.')
    main_chapter_id = path_components[0]

    current_level = chapter_content.get(main_chapter_id, {})

    # 对于路径中的每个子部分，向下导航到相应子章节
    curr_path = f"{path_components[0]}"
    init_depth = 1

    for i in range(init_depth, len(path_components)):
        curr_path += f".{path_components[i]}"
        if "children" in current_level and curr_path in current_level["children"]:
            current_level = current_level["children"][curr_path]
        else:
            # 找不到路径中的子章节
            return f"章节 {chapter_path} 在文档中不存在"

    return current_level


def replace_code(match, codes_info):
    # 获取匹配的代码段编号（re.sub(r"【代码段(\d+)】"）
    code_index = match.group(1)
    if code_index in codes_info:
        desc = codes_info[code_index].get("description", "")
        codes = codes_info[code_index].get("codes", "")
        replacement = f"代码描述：{desc}; 代码：\n```cpp\n{codes}\n```"
        return replacement
    else:
        # 若没有找到则返回原匹配字符串
        return match.group(0)


def substitute_codes(current_level, chapter_code_segments, is_chapter_16):
    """
    递归地替换章节及其子章节中的代码段引用

    Args:
        current_level: 当前章节内容
        chapter_code_segments: 整个主章节中提取的所有代码段信息

    Returns:
        更新后的章节内容
    """
    # 替换当前章节中的代码段引用
    if "原文内容" in current_level and current_level["原文内容"]:
        processed_text = current_level["原文内容"]
        codes_info = current_level["代码段"] if is_chapter_16 else chapter_code_segments
        processed_text = re.sub(r"【代码段(\d+)】", lambda match: replace_code(match, codes_info), processed_text)
        current_level["原文内容"] = processed_text

    # 递归处理所有子章节
    if "children" in current_level:
        updated_children = {}
        for child_id, child_content in current_level["children"].items():
            updated_children[child_id] = substitute_codes(child_content, chapter_code_segments, is_chapter_16)
        current_level["children"] = updated_children

    return current_level


def json_content_to_str(current_level_content, current_path):
    """
    递归地将章节内容转换为字符串，按照章节路径顺序拼接

    Args:
        current_level_content: 当前章节内容
        current_path: 当前章节路径

    Returns:
        拼接好的字符串
    """
    if not current_level_content:
        return ""

    result = ""

    # 获取当前章节信息
    title = current_level_content.get("title", "")
    summary = current_level_content.get("summary", "")
    original_text = current_level_content.get("原文内容", "")

    # 添加当前章节的基本信息
    if current_path:
        result += f"章节 {current_path}: {title}\n\n"
        if summary:
            result += f"摘要: {summary}\n\n"
        if original_text:
            result += f"{original_text}\n\n"

    # 处理子章节，确保按顺序排列
    if "children" in current_level_content and current_level_content["children"]:
        # 获取所有子章节ID并排序
        children_ids = list(current_level_content["children"].keys())

        # 自然排序子章节ID (处理数字顺序，例如10应该在2之后)
        def natural_sort_key(s):
            import re
            return [int(text) if text.isdigit() else text.lower()
                    for text in re.split(r'(\d+)', s)]

        children_ids.sort(key=natural_sort_key)

        # 按顺序递归处理每个子章节
        for child_id in children_ids:
            child_content = current_level_content["children"][child_id]
            # 使用完整路径处理子章节
            child_path = child_id
            result += json_content_to_str(child_content, child_path)

    return result


def load_specific_chapter_with_codes(chapter_id, chapters_dir):
    # 加载主章节内容
    chapter_content, chapter_code_segments, is_chapter_16 = load_full_chapter_content(chapter_id, chapters_dir)

    if chapter_content:
        # 只提取特定章节路径的内容
        current_level_content = extract_specific_chapter_text(chapter_content, chapter_id)
        current_level_content = substitute_codes(current_level_content, chapter_code_segments, is_chapter_16)
        specific_chapter_text = json_content_to_str(current_level_content, chapter_id)
    else:
        specific_chapter_text = None

    return specific_chapter_text


def load_doc_structure(doc_structure_file):
    doc_structure = load_document_structure(doc_structure_file)
    sections_info = extract_sections_info(doc_structure)
    return sections_info


def read_sample_code(file_or_folder_path:str):
    """从文件中读取样例Ascend C代码"""
    sample_code = ""
    
    try:
        if os.path.isdir(file_or_folder_path):
            for file_name in os.listdir(file_or_folder_path):
                if file_name.endswith((".md",".markdown")):  # 只读取 Markdown 文件
                    file_path = os.path.join(file_or_folder_path, file_name)
                    with open(file_path, "r", encoding="utf-8") as file:
                        content = file.read()
                        sample_code += content
        else:
            with open(file_or_folder_path, 'r', encoding='utf-8') as file:
                sample_code = file.read()
        return sample_code
    except Exception as e:
        print(f"读取样例代码文件时出错: {e}")
        return None

def hard_coded_chapter_ids(related_chapter_ids:list,hard_code_cpt):
    related_chapter_ids.extend(hard_code_cpt)
    unique_list = list(dict.fromkeys(related_chapter_ids))
    return unique_list

def match_chapter_ids_in_response(related_sections_response: str, task_id: int):
    # 提取章节编号，使用正则表达式匹配\chap{}格式
    related_chapter_ids = []
    chap_pattern = r'\\chap\{([0-9.]+)\}'
    matches = re.findall(chap_pattern, related_sections_response)

    if matches:
        related_chapter_ids = matches
    else:
        # 如果没有找到\chap{}格式，尝试使用备选方法提取
        raise RuntimeError(f"警告: 问题 {task_id} 未找到\\chap格式的章节引用")

    return related_chapter_ids


def load_chapter_main(related_chapter_ids, chapters_dir):
    chapter_contents = []
    loaded_chapters = set()  # 使用集合跟踪已加载的章节路径，避免重复

    for chapter_id in related_chapter_ids:
        if chapter_id in loaded_chapters:
            continue
        loaded_chapters.add(chapter_id)

        specific_chapter_text = load_specific_chapter_with_codes(chapter_id, chapters_dir)

        if specific_chapter_text:
            chapter_contents.append(specific_chapter_text)

    all_chapter_contents = "\n\n---\n\n".join(chapter_contents)
    return all_chapter_contents

